#########################################################################
#				checkpointutils.py : Checkpoint Utilities				#
#																		#
#########################################################################

#!/usr/bin/env python

import os, re, subprocess as sb
import time
import cnf   
import ss
import log

# Returns the loop dev associated to a particular uuid's nilfs2 image
base = '/vm/sslogs/'

def GetLoopDev(vm):
	
	# Get the name of the VM
	uuid = vm.domainCfgVars['uuid']
	# Get the loop device
	cmd1='/sbin/losetup -a'
	cmd2="awk -F\':\' \'/" + uuid + "/ {print $1}\'"
	p_loop = sb.Popen('%s' % cmd1, shell=True, stdout=sb.PIPE)
	p_awk = sb.Popen('%s' % cmd2, shell=True,stdin=p_loop.stdout,stdout=sb.PIPE)
	loopdev = p_awk.communicate()[0]
	
	# Return
	return loopdev

# Performs the xen save and nilfs snapshot
# TODO : 1) Code shell script functionality in python
#	 2) Implement versioning
#	 3) Currently implements name and not UUID. Add UUID support as well as name
#	 4) Implement logging
########################## WORKS ###########################################
def PerformSaveSnapshot(vm):

	loopdev = GetLoopDev(vm)
	name = str(vm.domainCfgVars['name'])
	uuid = str(vm.domainCfgVars['uuid'])
	
	version = str(GetVersion(uuid))
	
	log.write( 'Creating snapshot of domain ' + name + 'version number :- ' + version)

	#Execute xm save <name> <checkpoint file>
	savePath = '/vm/checkpoints/' + uuid
	localPath = '/tmp/' + uuid + '_' + version

	if not os.path.exists(savePath):
                os.system('mkdir ' + savePath)
		log.write( 'Created ' + savePath)
	
	save_cmd = '/usr/sbin/xm save ' + name + ' ' + localPath 
	p_save = sb.Popen('%s' %save_cmd, shell=True)
	#p_save.wait()
	log.write( 'xm save checkpoint saved')

	#Create Nilfs Snapshot	
	nilfs_cmd = 'mkcp -s ' + str(loopdev)
	p_nilfs = sb.Popen('%s' %nilfs_cmd, shell=True)
	#p_nilfs.wait()
	log.write( 'Nilfs snapshot saved')

	#Wait for save to complete
	p_save.wait()
	p_nilfs.wait()
	
	#Restore VM
	restore_cmd = '/usr/sbin/xm restore ' + localPath
	p_restore = sb.Popen('%s' %restore_cmd, shell=True)
	#p_restore.wait()
	log.write( 'Domain restored')

	#mv local file to NAS
	mv_cmd = 'mv ' + localPath + ' ' + savePath + '/chk_' + version
	p_mv = sb.Popen('%s' %mv_cmd, shell=True,stdout=sb.PIPE,stderr=sb.PIPE)
	#p_mv.wait()
	log.write( 'Moving local checkpoint to ' + savePath)
	
	# Create snapshot object
	ssObj = CreateSnapshotObject(vm.domainCfgVars['uuid'],loopdev)
	
	# Write to file
	WriteSS(ssObj)
	log.write( 'Snapshot information written to file')
	
	# Return snapshot object
	return ssObj
        
# Performs nilfs setup of a VM's disk image and copies the image from <path> to NAS
# Typically used to initialize a new VM. Flag = 0 -> Initial setup. Flag = 1 -> migration setup

# TODO : 1) See if Alan's UnionFS idea can be incorporated
#	 2) Implement boolean flag which toggles copying the image to NAS
#	 3) If flag exists, then modify to support associating loop dev on new host 
#	    after cold migration
#	 4) TEST!
#	 5) Implement exception handling
#	 6) Is this the right place for this function?
def CreateNilfsPartition(uuid,version):
	
	os.system('modprobe nilfs2');
	logfilePath = '/vm/nilfsimg/' + uuid + '/logfile_'+ str(uuid) + '_' + str(version) + '.img'
	
	# Creating sparse nilfs2 logfile image
	os.system('dd of='+ logfilePath +' bs=1G count=0 seek=12')
	
	# Creating nilfs2 filesystem on sparse file
	os.system('/sbin/mkfs.nilfs2 ' + logfilePath)
	
	return logfilePath

#Mounts the associated loopdev onto a mount path
def MountNilfs(loopdev, mntPath):

	if not os.path.exists(mntPath):
                os.system('mkdir ' + mntPath)
        cmd3 = '/sbin/mount.nilfs2 %s%s' %(loopdev,mntPath)
        p_mount = sb.Popen('%s' %cmd3, shell=True)
        p_mount.wait()

#Sets up the first available loop device to the given logfile
def NilfsLoopSetup(path):

	cmd = '/sbin/losetup -f'
        p_loopdev = sb.Popen ('%s' % cmd, shell=True,stdout=sb.PIPE)
        loopdev = p_loopdev.communicate()[0]
        loopdev = re.sub("\s+",' ',loopdev)
        
        cmd1 = '/sbin/losetup ' + loopdev +  path
        p_loopdev = sb.Popen('%s' % cmd1, shell=True)
        p_loopdev.wait()
	return loopdev

#Performs cold migration (restore) of a VM, given the uuid and snapshot version.
def Migrate(vm):
	
	uuid = str(vm.domainCfgVars['uuid'])
	log.write( 'Migrating ' + vm.domainCfgVars['name'])

	# Find the name of the current disk image
	p_old = sb.Popen('ls /vm/nilfsimg/' + uuid + '/logfile*', shell=True, stdout=sb.PIPE)
	oldLogFile = str(p_old.communicate()[0])
	#oldLogFile = '/vm/nilfsimg/' + uuid + '/' + oldLogFileName	
	log.write( 'Current (out of sync) image is ' + oldLogFile)
	
	# Setup the current disk image to a loop device
	oldLoopDev = NilfsLoopSetup(oldLogFile)
	log.write( 'Old loop dev setup with ' + str(oldLoopDev))
	
	# Mount the current disk image onto a mountpoint (used to extract snapshot)
	tmpPath = '/vm/tmp/' + uuid	
	if not os.path.exists(tmpPath):
		os.system('mkdir ' + tmpPath)
	
	tmpMntPath = tmpPath +  '/logmnt'
	if not os.path.exists(tmpMntPath):
		os.system('mkdir ' + tmpMntPath)

	MountNilfs(oldLoopDev,tmpMntPath)
	log.write( str(oldLoopDev) + ' mounted to ' + tmpMntPath)

	# Get the most recent snapshot
	lastSS = vm.GetLastSnapshot()
	if lastSS == None:
		return False
	
	log.write( 'Restore VM from version no ' + lastSS.version + ' cno='+  str(lastSS.cno))
	# Create a partition recovered disk image (recovered from snapshot)
	logfilePath = CreateNilfsPartition(uuid,lastSS.version)
	log.write( 'New nilfs partition created')
	
	# Setup the new logfile to a loop device
	loopdev = NilfsLoopSetup(logfilePath)
	log.write( 'New partition set up with ' + str(loopdev))

	# Mount new logfile to a mountpath
	mntPath = '/vm/mnt/' + uuid + '/'
	MountNilfs(loopdev,mntPath)
	log.write( 'Partition mounted at ' + mntPath)
	
	# Mount snapshot (cno) and copy data to new logfle
	SnapshotMountCopy(oldLoopDev,uuid,lastSS) 
	log.write( 'Snapshot mounted and copied to new partition')

	# Restore vm
	restore_cmd = '/usr/sbin/xm restore /vm/checkpoints/' + uuid + '/chk_' + str(lastSS.version)
	p_restore = sb.Popen('%s' %restore_cmd, shell=True)
	log.write( 'VM migrated')

	# Print
	log.write( 'Cleaning up')

	# Unmount the oldLogFile
	os.system('umount ' + tmpMntPath)
	log.write( tmpMntPath + ' unmounted')
	
	# Detach the old loopdev
	cmd = '/sbin/losetup -d ' + oldLoopDev
	p_detatch = sb.Popen('%s' %cmd, shell=True)
	p_detatch.wait()
	log.write( 'Loop device ' + str(oldLoopDev) + 'detached')
 
	# Remove the old log file		
	p_del_log = sb.Popen('rm -f %s' %oldLogFile,shell=True)
	log.write( 'Old logfile deleted')
	
	# Return
	return True
	

def SnapshotMountCopy(loopdev,uuid,s):

	#Check if paths exits
	snapshotMntPath = '/vm/tmp/' + str(uuid)
	if not os.path.exists(snapshotMntPath):
                os.system('mkdir ' + snapshotMntPath)
	snapshotMntPath += '/ssmnt/'
	if not os.path.exists(snapshotMntPath):
        	os.system('mkdir ' + snapshotMntPath)

	#Mount the snapshot to /vm/tmp/<uuid>/ssmnt
	cmd = '/sbin/mount.nilfs2 -r -o cp=' + str(s.cno) + ' ' + loopdev + '/vm/tmp/' + uuid + '/ssmnt'
	log.write( 'Snapshot mounted from cno=' + str(s.cno))

	p_ssmnt = sb.Popen('%s' %cmd, shell=True)
	p_ssmnt.wait()

	#Copy all data in snapshot to new location
	cmd2 = 'cp -rf /vm/tmp/' + uuid + '/ssmnt/* /vm/mnt/' + uuid + '/'
	p_cp = sb.Popen('%s' %cmd2, shell=True)
	p_cp.wait()
	log.write( 'All data from snapshot moved to new location')
	#Unmount snapshot
	cmd3 = 'umount ' + '/vm/tmp/' + uuid + '/ssmnt/'
	p_umount = sb.Popen('%s' %cmd3, shell=True)
	log.write( 'snapshot unmounted')
	
	#Remove all the old checkpoints and logs
	CleanCheckpointLogs(uuid)

# Creates a snapshot object
def CreateSnapshotObject(uuid,loopdev):
	
	s = ss.snapshot()
	s.SetTimestampNow()
	s.cno = GetCNO(loopdev)
	s.version = str(GetVersion(uuid))
	s.uuid = uuid
	return s
	
# Writes snapshot information to file	
def WriteSS(ss): 
	
	filename = base + ss.uuid
	if not os.path.exists(filename):
		os.system('mkdir ' + filename)
	filename += '/log.txt'
	WRITE = str(ss.version) + '|' + str(ss.cno) + '|' + ss.GetTimestampString()
	FILE = open(filename, "a")
    	FILE.write(WRITE + '\n')
    	log.write( 'Snapshot info version: ' + str(ss.version) + ' cno: ' + str(ss.cno) + 'written to file')

# Reads and lists snapshot info
def ListSS(uuid):

	# Build path and filename
	filePath = base + uuid
	fileName = filePath + '/log.txt'
	
	# Make sure the path exists
	if not os.path.exists(filePath) or not os.path.exists(fileName):
		log.write( 'No Snapshots Taken')
		return []
	
	# Open the file
	logFile = open(fileName, "r")
	
	# Read lines and create list
	snapshotList = []
	for line in logFile.readlines():
		
		# Read fields
		version = line.split('|')[0]
		cno = line.split('|')[1]
		timeStampString = line.split('|')[2].strip()
		
		# Create object
		sobj = ss.snapshot()
		sobj.version = version
		sobj.cno = cno
		sobj.uuid = uuid
		sobj.SetTimestampFromString(timeStampString)
		
		# Append
		snapshotList.append(sobj)
	
	# Return list
	return snapshotList

# Prints a list of snapshots for the uuid	
def PrintSSList(uuid):
	s = ListSS(uuid)
	for snap in s:
		log.write( '%7s\t%7s\t15%s' %(snap.version,snap.cno,snap.GetTimestampString()))
 

#Returns the cno number of the last snapshot performed
def GetCNO(loopdev):

	tail = 'tail -1'
	lscp = 'lscp -s ' + loopdev
	p_lscp = sb.Popen('%s' %lscp, shell=True, stdout = sb.PIPE)
	p_tail = sb.Popen('%s' %tail, shell=True, stdin=p_lscp.stdout, stdout = sb.PIPE)
	p_lscp.wait()
	cline = str(p_tail.communicate()[0])
	cno = re.sub("\s+","|",cline).split('|')
	return int(cno[1])

#Gets the last snapshot version and returns an auto incremented value
def GetVersion(uuid):
	
	filename = base + uuid
        if not os.path.exists(filename):
                return 1
        filename += '/log.txt'
	cmd = 'tail -1 ' + filename
	p_cmd = sb.Popen('%s' %cmd, shell=True, stdout=sb.PIPE)
	version = str(p_cmd.communicate()[0]).split('|')[0]	
	return int(version) + 1

#Cleans up all the old checkpoint and log data (in case of fast migrate)
def CleanCheckpointLogs(uuid):
	
	rm_cmd1 = 'rm -rf /vm/checkpoints/' + uuid + '/*'
	p_rm1 = sb.Popen('%s' %rm_cmd1, shell=True)
	
	rm_cmd2 = 'rm -rf /vm/sslogs/' + uuid + '/*'
	p_rm2 = sb.Popen('%s' %rm_cmd2, shell=True)
	
	log.write('Cleaned up checkpoints and logs for ' + uuid)


